package com.zyzg.filter;

import javax.servlet.*;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.util.concurrent.ConcurrentHashMap;

import com.alibaba.fastjson2.JSONObject;
import org.springframework.stereotype.Component;

@Component
public class IpRateLimitFilter implements Filter {

    // 存储每个 IP 的访问计数器
    private static final ConcurrentHashMap<String, Integer> ipRequestCountMap = new ConcurrentHashMap<>();
    // 存储每个 IP 首次访问的时间戳
    private static final ConcurrentHashMap<String, Long> ipFirstAccessTimeMap = new ConcurrentHashMap<>();
    // 时间窗口（一分钟，单位：毫秒）
    private static final long TIME_WINDOW = 60 * 1000;
    // 最大请求次数
    private static final int MAX_REQUESTS = 30;

    @Override
    public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain)
            throws IOException, ServletException {
        HttpServletRequest httpRequest = (HttpServletRequest) request;
        HttpServletResponse httpResponse = (HttpServletResponse) response;
        String clientIp = getClientIp(httpRequest);

        long currentTime = System.currentTimeMillis();
        // 检查该 IP 是否在时间窗口内
        if (ipFirstAccessTimeMap.containsKey(clientIp)) {
            long firstAccessTime = ipFirstAccessTimeMap.get(clientIp);
            if (currentTime - firstAccessTime > TIME_WINDOW) {
                // 超出时间窗口，重置计数器和时间戳
                ipRequestCountMap.put(clientIp, 1);
                ipFirstAccessTimeMap.put(clientIp, currentTime);
            } else {
                // 在时间窗口内，增加计数器
                int requestCount = ipRequestCountMap.get(clientIp);
                if (requestCount >= MAX_REQUESTS) {
                    // 超过最大请求次数，返回 429 状态码
                    httpResponse.setStatus(429);
                    JSONObject responseJson = new JSONObject();
                    responseJson.put("code", 429);
                    responseJson.put("message", "请求过于频繁，请稍后再试");
                    httpResponse.getWriter().write(responseJson.toJSONString());
                    return;
                }
                ipRequestCountMap.put(clientIp, requestCount + 1);
            }
        } else {
            // 首次访问，初始化计数器和时间戳
            ipRequestCountMap.put(clientIp, 1);
            ipFirstAccessTimeMap.put(clientIp, currentTime);
        }

        // 放行请求
        chain.doFilter(request, response);
    }

    private String getClientIp(HttpServletRequest request) {
        String xffHeader = request.getHeader("X-Forwarded-For");
        if (xffHeader == null) {
            return request.getRemoteAddr();
        }
        return xffHeader.split(",")[0];
    }
}